"""
Task-specific callbacks for continual learning.

This module provides callback classes for handling task-level events in continual learning
scenarios, extending PyTorch Lightning's callback system.
"""

from typing import Any, List, Optional
from pytorch_lightning.callbacks import Callback
import pytorch_lightning as pl
from torch.utils.data import DataLoader


class TaskCallback(Callback):
    """
    Base callback for task-level events in continual learning.

    This callback provides hooks for task transitions, which are not
    available in standard PyTorch Lightning callbacks. Other callbacks
    can inherit from this to provide task-aware functionality.

    Attributes:
        current_task (int): The current task being processed
    """

    def __init__(self):
        super().__init__()
        self.current_task = 0

    def on_task_start(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", task_id: int
    ) -> None:
        """
        Called when a new task begins.

        Args:
            trainer: PyTorch Lightning trainer
            pl_module: PyTorch Lightning module
            task_id: Identifier of the new task
        """
        self.current_task = task_id

    def on_task_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        task_id: int,
        dataloader: Optional[DataLoader] = None,
    ) -> None:
        """
        Called when a task is completed.

        Args:
            trainer: PyTorch Lightning trainer
            pl_module: PyTorch Lightning module
            task_id: Identifier of the completed task
            dataloader: Optional dataloader with the task's data
        """
        pass


class ComponentTaskCallback(TaskCallback):
    """
    Callback that bridges PyTorch Lightning callbacks with component system.

    This callback wraps BaseComponents and calls their task-level hooks at
    appropriate points in the training process.

    Args:
        components: List of continual learning components to manage
    """

    def __init__(self, components: List[BaseComponent]):
        super().__init__()
        self.components = components

    def on_task_start(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", task_id: int
    ) -> None:
        """Notify all components of task start."""
        super().on_task_start(trainer, pl_module, task_id)

        for component in self.components:
            component.on_task_start(task_id)

    def on_task_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        task_id: int,
        dataloader: Optional[DataLoader] = None,
    ) -> None:
        """Notify all components of task end."""
        super().on_task_end(trainer, pl_module, task_id, dataloader)

        for component in self.components:
            component.on_task_end(task_id, dataloader)

    def on_train_epoch_start(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Notify all components of epoch start."""
        for component in self.components:
            component.on_epoch_start()

    def on_train_epoch_end(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Notify all components of epoch end."""
        for component in self.components:
            component.on_epoch_end()

    def on_train_batch_start(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        batch: Any,
        batch_idx: int,
        unused: Optional[int] = 0,
    ) -> None:
        """Notify all components of batch start."""
        for component in self.components:
            component.on_batch_start(batch, batch_idx)

    def on_train_batch_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        outputs: Any,
        batch: Any,
        batch_idx: int,
        unused: Optional[int] = 0,
    ) -> None:
        """Notify all components of batch end."""
        for component in self.components:
            component.on_batch_end(batch, batch_idx, outputs)


class EvaluationTaskCallback(TaskCallback):
    """
    Callback for task-specific evaluation in continual learning.

    This callback collects per-task metrics and logs them appropriately.

    Attributes:
        metrics (Dict): Dictionary to store metrics per task
    """

    def __init__(self):
        super().__init__()
        self.metrics = {}

    def on_task_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        task_id: int,
        dataloader: Optional[DataLoader] = None,
    ) -> None:
        """
        Perform evaluation at the end of a task.

        If previous tasks are available, this will trigger evaluation on them
        to measure forgetting.
        """
        super().on_task_end(trainer, pl_module, task_id, dataloader)

        # Store current task metrics
        if (
            hasattr(pl_module, "eval_component")
            and pl_module.eval_component is not None
        ):
            metrics = pl_module.eval_component.compute_metrics()
            self.metrics[task_id] = metrics

            # Log overall metrics
            for name, value in metrics.items():
                trainer.logger.log_metrics({f"task_{task_id}/{name}": value})


class CheckpointTaskCallback(TaskCallback):
    """
    Callback for saving checkpoints at task boundaries.

    Args:
        checkpoint_dir: Directory to save checkpoints
        save_last_only: Whether to keep only the most recent task checkpoint
    """

    def __init__(
        self, checkpoint_dir: str = "checkpoints", save_last_only: bool = False
    ):
        super().__init__()
        self.checkpoint_dir = checkpoint_dir
        self.save_last_only = save_last_only

        # Create checkpoint directory
        import os

        os.makedirs(checkpoint_dir, exist_ok=True)

    def on_task_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        task_id: int,
        dataloader: Optional[DataLoader] = None,
    ) -> None:
        """Save a task-specific checkpoint when a task is completed."""
        super().on_task_end(trainer, pl_module, task_id, dataloader)

        # Save task-specific checkpoint
        trainer.save_checkpoint(f"{self.checkpoint_dir}/task_{task_id}_checkpoint.ckpt")

        # If save_last_only and not the first task, remove previous checkpoint
        if self.save_last_only and task_id > 0:
            import os

            prev_checkpoint = (
                f"{self.checkpoint_dir}/task_{task_id - 1}_checkpoint.ckpt"
            )
            if os.path.exists(prev_checkpoint):
                os.remove(prev_checkpoint)
